import os
import sys

root_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../'))
sys.path.append(root_path)

from shell.knowledge_graph.text_classification.predict import *


class Classifier(object):
    def __init__(self):
        model_path = root_path + "/shell/knowledge_graph/text_classification/"
        self.parent_model = load_model(model_path + 'parent_model.pkl', True)
        self.child_model = load_model(model_path + 'child_model.pkl', False)
        self.parent_dic = {}
        with open(model_path + 'parent_label.txt', encoding='utf-8')as file:
            for line in file.readlines():
                k, v = line.strip().split(' ')
                self.parent_dic[int(v)] = k
        self.child_dic = {}
        with open(model_path + 'child_label.txt', encoding='utf-8')as file:
            for line in file.readlines():
                k, v = line.strip().split(' ')
                self.child_dic[int(v)] = k

    def predict(self, ques):
        index = seq2index(ques)
        index = np.expand_dims(index, 0)
        x = padding_seq(index)
        x = torch.from_numpy(x)
        if torch.cuda.is_available():
            x = x.cuda()
        y_p = self.parent_model(x)
        y_c = self.child_model(x)
        y_p = torch.argmax(y_p[0]).cpu().item()
        y_c = torch.argmax(y_c[0]).cpu().item()
        return self.parent_dic[y_p], self.child_dic[y_c]


if __name__ == '__main__':
    c = Classifier()
    a, b = c.predict('236822199504108770怎么担保的')
    print('主节点：', a, '子节点：', b)
    a = c.predict('5574220148709795的公司地址？')
    print('主节点：', a[0], '子节点：', a[1])
